package com.koicarp.agent.provider;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.List;


import com.koicarp.agent.common.constant.CommonConstant;
import com.koicarp.agent.provider.constant.ProviderErrMsgConstant;
import com.koicarp.agent.provider.entity.EmbedDbModel;
import dev.langchain4j.store.embedding.pgvector.MetadataStorageConfig;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import lombok.extern.slf4j.Slf4j;

/**
 * @ClassName:  CustomPgVectorEmbeddingStore   
 * @Description: 自定义向量存储
 * @author: liutao
 * @date:   2025年3月10日 上午9:21:20    
 */
@Slf4j
public class CustomPgVectorEmbeddingStore extends PgVectorEmbeddingStore {

	public CustomPgVectorEmbeddingStore(String host,
            Integer port,
            String user,
            String password,
            String database,
            String table,
            Integer dimension,
            Boolean useIndex,
            Integer indexListSize,
            Boolean createTable,
            Boolean dropTableFirst,
            MetadataStorageConfig metadataStorageConfig) {
		super(host, port, user, password, database, table, dimension, 
				useIndex, indexListSize, createTable, dropTableFirst, metadataStorageConfig);
	}
	
	private static final String EMBEDDING_ID = "embedding_id";
	
//	private static final String EMBEDDING = "embedding";
	
	private static final String TEXT = "text";
	
	private static final String METADATA = "metadata";
	
	
	public List<EmbedDbModel> searchByVectorIds(List<String> vectorIds)  {
		List<EmbedDbModel> models = new ArrayList<>();
		
		StringBuilder builder = new StringBuilder();
		for (int i = 0; i < vectorIds.size(); i++) {
			if (vectorIds.get(i) != null) {
				builder.append("'").append(vectorIds.get(i)).append("'").append(CommonConstant.COMMAN);
			}
		}
		String inSql = builder.substring(0, builder.length() - 1);
		
		try (Connection connection = getConnection()) {
			 String query = String.format("select * from %s c where c.embedding_id  in (%s)",
					super.table, inSql);
			 try (PreparedStatement selectStmt = connection.prepareStatement(query)) {
	                try (ResultSet resultSet = selectStmt.executeQuery()) {
	                    while (resultSet.next()) {
	                    	EmbedDbModel model = EmbedDbModel.builder()
	                    	.vectorId(resultSet.getString(EMBEDDING_ID))
	                    	.text(resultSet.getString(TEXT))
	                    	.metadata(resultSet.getString(METADATA))
	                    	.build();

	                        // 暂不处理向量字段，需要时放开
	                        // PGvector vector = (PGvector) resultSet.getObject("embedding");
	                        // Embedding embedding = new Embedding(vector.toArray());
	                        models.add(model);
	                    }
	                }
	            }
		} catch (Exception e) {
			log.error(ProviderErrMsgConstant.EMBED_TEXT_ERROR, super.table, inSql, e.getMessage());
		}
		return models;
	}
}
